"""
Copyright (c) 2022 Guangdong University of Technology
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Meng Xiang, Junjiang Xiang
Created: 2023/8/19
Supported by: National Key Research and Development Program of China
"""

import commpy as cp
import numpy as np
from scipy.signal import lfilter


def pulse_shaper(input, up_sampling_factor, roll_off, total_baud):
    """RRC 滤波器
    Args:
        input[0]: 输入X偏振信号, numpy类型
        input[1]: 输入Y偏振信号, numpy类型
        up_sampling_factor: 上采样倍数
        roll_off:滚降因子
        baudrate:波特率

    Returns:
        output[0]：输出X偏振信号
        output[1]:输出Y偏振信号
    """
    baudrate = total_baud
    X_Pol = input[0]
    Y_Pol = input[1]

    if np.ndim(X_Pol) == 1:
        time_idx, rrc_filter = cp.rrcosfilter(N=256 * up_sampling_factor, alpha=roll_off, Ts=1 / baudrate,
                                              Fs=up_sampling_factor * baudrate)
        X_Pol = lfilter(rrc_filter, 1, X_Pol, axis=0)
        Y_Pol = lfilter(rrc_filter, 1, Y_Pol, axis=0)
    else:
        X_Pol = np.squeeze(X_Pol, axis=1)
        Y_Pol = np.squeeze(Y_Pol, axis=1)
        time_idx2, rrc_filter2 = cp.rrcosfilter(N=256 * up_sampling_factor, alpha=roll_off, Ts=1 / baudrate,
                                                Fs=up_sampling_factor * baudrate)
        X_Pol = lfilter(rrc_filter2, 1, X_Pol, axis=0)
        Y_Pol = lfilter(rrc_filter2, 1, Y_Pol, axis=0)
        X_Pol = np.expand_dims(X_Pol, axis=1)
        Y_Pol = np.expand_dims(Y_Pol, axis=1)

    
    output = [X_Pol, Y_Pol]

    return output



